"""
Copyright (c) 2022 Guangdong University of Technology
PhotLab is licensed under [Open Source License].
You can use this software according to the terms and conditions of the [Open Source License].
You may obtain a copy of [Open Source License] at: [https://open.source.license/]

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.

See the [Open Source License] for more details.

Author: Meng Xiang, Junjiang Xiang
Created: 2023/8/19
Supported by: National Key Research and Development Program of China
"""

import numpy as np


def circshift(u, shift_index):
    h = len(u)
    if shift_index < 0:
        u = np.vstack((u[-shift_index:, :], u[:-shift_index, :]))
    else:
        u = np.vstack((u[(h-shift_index):, :], u[:(h-shift_index), :]))
    return u


def SP_SymbolTimingEstimatorSingleBlock(rollOffFactor, Rxdata_2SaPerSym):
    NBlocks = 1
    g = 2 * np.sin(np.pi * rollOffFactor / 2) / (np.pi * (4 - rollOffFactor ** 2)) / (rollOffFactor / 4)
    estimatedTimingOffset = np.zeros((1, NBlocks))

    for jj in range(NBlocks):
        env1 = abs(Rxdata_2SaPerSym) ** 2
        env2 = np.real(np.conj(Rxdata_2SaPerSym) * circshift(Rxdata_2SaPerSym, -1))
        n = np.arange(0, (len(Rxdata_2SaPerSym)))

        estimatedTimingOffset[:, jj] = 1/(2 * np.pi) * np.angle(g * np.sum(env1[:, 0] * np.exp(-1j * n * np.pi)) + sum(env2[:-1,0] * np.exp(-1j * np.pi*(n[:-1]-0.5))))

    return estimatedTimingOffset


def time_recovery(input, FilterRollOff = 0.1, BlockSize = 3000):
    """

    Args:
        input: 信号
        FilterRollOff: RRC滚降因子
        BlockSize: 窗口大小

    Returns:

    """
    SamplesPerSymbol = 2
    NBlocks = int(np.ceil(len(input) / BlockSize))

    BLK_addresses = np.arange(1, (NBlocks) * BlockSize, BlockSize)

    L = len(input)
    output = np.zeros((1, int(L)), dtype=np.complex)
    SFO_Vector = np.zeros((1, NBlocks + 1))

    for jj in range(0, NBlocks):
        if jj < NBlocks-1:
            BLK_range = np.arange(BLK_addresses[jj], (BLK_addresses[jj] + BlockSize))
        else:
            BLK_range = np.arange(BLK_addresses[jj], L+1)

        print(jj)
        print(BLK_range)
        sfo = SP_SymbolTimingEstimatorSingleBlock(FilterRollOff, input[BLK_range-1, :])

        dt = sfo * SamplesPerSymbol

        SFO_Vector[:, jj + 1] = SFO_Vector[:, jj] + dt

        omega = np.hstack((np.arange(0, (np.ceil(L / 2))), np.arange(np.ceil(-L / 2), 0)))

        X2 = np.fft.fft(input.T)
        X1_DELAY = X2 * np.exp(2j * np.pi * omega * dt / L)
        x1_delay = np.fft.ifft(X1_DELAY)

        output[:, BLK_range-1] = x1_delay[:, BLK_range-1]
        input = x1_delay.T

    return output

